package com.artislong.servlet;

import cn.hutool.extra.servlet.JakartaServletUtil;
import cn.hutool.http.ContentType;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import com.fasterxml.jackson.databind.ObjectMapper;
import io.modelcontextprotocol.client.McpClient;
import io.modelcontextprotocol.client.McpSyncClient;
import io.modelcontextprotocol.client.transport.ServerParameters;
import io.modelcontextprotocol.client.transport.StdioClientTransport;
import io.modelcontextprotocol.spec.McpSchema;
import jakarta.servlet.ServletConfig;
import jakarta.servlet.ServletException;
import jakarta.servlet.annotation.WebServlet;
import jakarta.servlet.http.HttpServlet;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;

import java.io.IOException;
import java.time.Duration;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

/**
 * @author 陈敏
 * @version McpClientServlet.java, v 1.0 2025 07 09 16:54 chenmin Exp $
 * Created on 2025 07 09
 */
@WebServlet(asyncSupported = true)
public class McpClientServlet extends HttpServlet {

    private McpSyncClient mcpSyncClient;
    private final ObjectMapper OBJECT_MAPPER = new ObjectMapper();

    @Override
    public void init(ServletConfig config) throws ServletException {
        super.init(config);
        StdioClientTransport stdioClientTransport = new StdioClientTransport(
                ServerParameters.builder("java")
                        .args("-jar", "G:/BigModel/code/mcp-sample/mcp-server/mcp-server-stdio/target/mcp-server-stdio-1.0.jar")
                        .build(),
                OBJECT_MAPPER
        );

        McpSchema.ClientCapabilities capabilities = McpSchema.ClientCapabilities.builder()
                .roots(true)      // Enable filesystem roots support with list changes notifications
                .sampling()       // Enable LLM sampling support
                .build();

        // Create a sync client with custom configuration
        mcpSyncClient = McpClient.sync(stdioClientTransport)
                .clientInfo(new McpSchema.Implementation("my-client", "1.0.0"))
                .initializationTimeout(Duration.ofSeconds(1000))
                .requestTimeout(Duration.ofSeconds(1000))
                .capabilities(capabilities)
                .build();

        // Initialize connection
        mcpSyncClient.initialize();
        mcpSyncClient.setLoggingLevel(McpSchema.LoggingLevel.INFO);
    }

    @Override
    protected void doGet(HttpServletRequest req, HttpServletResponse resp) throws ServletException, IOException {
        resp.setCharacterEncoding("UTF-8");
        resp.setContentType("application/json");
        String query = req.getParameter("query");
        switch (query) {
            case "tools":
                McpSchema.ListToolsResult tools = mcpSyncClient.listTools();
                List<JSONObject> toolMetaDatas = tools.tools().stream().map(tool -> {
                    JSONObject result = new JSONObject();
                    result.set(tool.name(), tool.description());
                    result.set("schema", JSONUtil.parse(com.alibaba.fastjson2.JSONObject.toJSONString(tool.inputSchema())));
                    return result;
                }).collect(Collectors.toList());
                JakartaServletUtil.write(resp, JSONUtil.toJsonStr(toolMetaDatas), ContentType.JSON.getValue());
                break;
            case "prompts":
                McpSchema.ListPromptsResult prompts = mcpSyncClient.listPrompts();
                List<JSONObject> promptMetaDatas = prompts.prompts().stream().map(prompt -> {
                    JSONObject result = new JSONObject();
                    result.set(prompt.name(), prompt.description());
                    result.set("arguments", prompt.arguments());
                    return result;
                }).collect(Collectors.toList());
                JakartaServletUtil.write(resp, JSONUtil.toJsonStr(promptMetaDatas), ContentType.JSON.getValue());
                break;
            case "resources":
                McpSchema.ListResourcesResult resources = mcpSyncClient.listResources();
                List<JSONObject> resourceMetaDatas = resources.resources().stream().map(resource -> {
                    JSONObject result = new JSONObject();
                    result.set(resource.name(), resource.description());
                    result.set("mimeType", resource.mimeType());
                    result.set("uri", resource.uri());
                    result.set("annotations", resource.annotations());
                    return result;
                }).collect(Collectors.toList());
                JakartaServletUtil.write(resp, JSONUtil.toJsonStr(resourceMetaDatas), ContentType.JSON.getValue());
                break;
            case "call":
                String method = req.getParameter("method");
                Map<String, Object> paramMap = new HashMap<>(JakartaServletUtil.getParamMap(req));
                paramMap.remove("method");
                paramMap.remove("query");
                McpSchema.CallToolResult result = mcpSyncClient.callTool(
                        new McpSchema.CallToolRequest(method, paramMap)
                );
                List<McpSchema.Content> content = result.content();
                JakartaServletUtil.write(resp, JSONUtil.toJsonStr(com.alibaba.fastjson2.JSONObject.toJSONString(content)), ContentType.JSON.getValue());
                break;
        }
    }

    @Override
    public void destroy() {
        // 关闭MCP服务器
        if (mcpSyncClient != null) {
            try {
                mcpSyncClient.close();
            } catch (Exception e) {
                getServletContext().log("Error closing MCP server", e);
            }
        }
    }
}